#include <string>
#include <algorithm>
#include <numeric>
#include <istream>
#include <vector>
#include <fstream>
#include <deque>
#include <map>
#include <stdexcept>
#include <ext/hash_map>


// boost functions
#include <boost/lexical_cast.hpp>

// my functions
#include <single_linkage.h>
#include <portable_timer.h>
#include <progress_indicator.h>
#include <print_error.h>
#include <ultoa.h>
#include <index_string.h>

#include "get_arg.h"

#include <gc.h>
using std::cerr;
using std::ifstream;
using std::ofstream;
using std::cout;
using std::string;
using std::vector;
using std::deque;
using std::sort;
using boost::lexical_cast;
using std::ostringstream;





typedef std::pair<unsigned, unsigned> t_id_pair;

struct t_clustered_id_pair : public std::pair<unsigned, unsigned>
{
	unsigned cluster;
	t_clustered_id_pair(unsigned id1_, unsigned id2_)
		:std::pair<unsigned, unsigned>(id1_, id2_)
	{
    }
};
operator < (const t_clustered_id_pair& a, const t_clustered_id_pair& b)
{
	return a.cluster< b.cluster;
}

//________________________________________________________________________________________

//		lookup_prot_name

//________________________________________________________________________________________
inline unsigned lookup_prot_name(t_index_string& prot_ids, const string& prot_name)
{
	unsigned prot_id;
	if (!prot_ids.is_indexed(prot_name, prot_id))
		throw std::runtime_error("Protein identifier in blast result [" + prot_name +
								"] not found in the supplied protein id to gene id "
								"lookup table.");
	return prot_id;
}







//________________________________________________________________________________________

//		parse_blast_pairs_from_file

//________________________________________________________________________________________
void parse_blast_pairs_from_file(	deque<unsigned>&					prot_id_to_gene_id,
									t_index_string&						prot_ids,
									t_index_string&						gene_ids,
									deque<t_id_pair>&					gene_pairs,
									deque<t_clustered_id_pair>&			prot_pairs,
									const string&						file_name,
									t_progress_indicator& 				dots)
{
	ifstream ifs(file_name.c_str());
	if (!ifs)
		throw std::runtime_error("Could not open the blast pairs input file [" + file_name + "]");

	dots.set_every(100000);

	unsigned cnt_lines = 0;

	// get data for one protein pair
	string prot_name1, prot_name2;
	ifs >> prot_name1 >> prot_name2;

	while (ifs.good())
	{
		cnt_lines++;
		dots++;

		// get protein and gene indices
		unsigned prot_id1 = lookup_prot_name(prot_ids, prot_name1);
		unsigned prot_id2 = lookup_prot_name(prot_ids, prot_name2);
		unsigned gene_id1 = prot_id_to_gene_id[prot_id1];
		unsigned gene_id2 = prot_id_to_gene_id[prot_id2];

		// ignore comparisons between transcripts of the same gene
		assert(gene_id1 != gene_id2);
		if (gene_id1 != gene_id2)
		{
			// sort gene indices so that comparisons don't have to be symmetrical
			gene_pairs.push_back(t_id_pair(gene_id1,  gene_id2));
			prot_pairs.push_back(t_clustered_id_pair(prot_id1,  prot_id2));
		}


		// get next line
		ifs >> prot_name1 >> prot_name2;
	}

	if (ifs.eof())
		return;
	throw std::runtime_error("Invalid blast results format on line #" +
								ultoa(cnt_lines) + " of " + file_name);
}


//________________________________________________________________________________________

//		parse_protid_gene_id_from_file

//________________________________________________________________________________________
void parse_protid_gene_id_from_file(const string&						file_name,
									deque<unsigned>&					prot_id_to_gene_id,
									t_index_string&						prot_ids,
									t_index_string&						gene_ids,
									t_progress_indicator& 				dots,
									t_program_args&						args)
{
	ifstream is(file_name.c_str());
	if (!is)
		throw std::runtime_error("Could not open the gene to protein identifier map file [" + file_name + "]");

	unsigned cnt_lines = 0;

	// get data for one protein pair
	string prot_name, gene_name;
	is >> prot_name >> gene_name;
	dots.set_every(5000);
	while (is)
	{
		dots++;
		cnt_lines++;
		bool added;
		unsigned prot_id = prot_ids.index(prot_name, added);
		lookup_prot_name(prot_ids, prot_name);
		// warn duplicatd gene_ids.string_at(prot_id_to_gene_id[prot_id]) and gene_name
		if (added)
		{
			unsigned gene_id = gene_ids.index(gene_name, added);
			// save gene_id lookup
			prot_id_to_gene_id.push_back(gene_id);

		}
		else
		{
			args.CERR() << "WARNING:\tDuplicated prot ID " << prot_name
						<< " with previous gene of " << gene_ids.string_at(prot_id_to_gene_id[prot_id])
						<< " and now with gene of " << gene_name << "\n";
		}

		// get next line
		is >> prot_name >> gene_name;
	}

	if (is.eof())
		return;
	throw std::runtime_error("Invalid protein / gene name on line #" +
								ultoa(cnt_lines) + " of " + file_name);


}



void create_gene_ids_to_cluster_lookup(	const std::deque<std::vector<unsigned> >& 	clusters,
										vector<unsigned>& 							gene_ids_to_cluster_lookup)
{
	for (unsigned ii = 0; ii < clusters.size(); ++ii)
		for (unsigned jj = 0; jj < clusters[ii].size(); ++jj)
			gene_ids_to_cluster_lookup[clusters[ii][jj]] = ii;
}

void assign_prot_pairs_to_clusters(		deque<t_clustered_id_pair>&	prot_pairs,
										const deque<unsigned>&		prot_id_to_gene_id,
										const vector<unsigned>& 	gene_ids_to_cluster_lookup)
{
	for (unsigned ii = 0; ii < prot_pairs.size(); ++ii)
	{
		prot_pairs[ii].cluster = gene_ids_to_cluster_lookup[prot_id_to_gene_id[prot_pairs[ii].first]];
	}
}

void print_cluster_sizes(	const std::deque<std::vector<unsigned> >&	clusters,
							const t_index_string&						gene_ids,
							t_program_args& 							args)
{
	string file_name = args.sizes;
	ofstream ofs(file_name.c_str());
	if (!ofs)
		throw std::runtime_error("Could not open the file [" + file_name +
									"] to print out cluster sizes");
	file_name = args.clusters;
	ofstream ofscluster(file_name.c_str());
	if (!ofscluster)
		throw std::runtime_error("Could not open the file [" + file_name +
									"] to print out gene clusters");

	std::map<unsigned, unsigned> cluster_sizes;
	for (unsigned i = 0; i < clusters.size(); ++i)
	{
		unsigned cluster_size = clusters[i].size();
		if (cluster_sizes.count(cluster_size))
			++cluster_sizes[cluster_size];
		else
			cluster_sizes[cluster_size] = 1;

		ofscluster << ">cluster #" << (i+1) << "\n";
		for (unsigned j = 0; j < cluster_size; ++j)
		{
			ofscluster << gene_ids.string_at(clusters[i][j]) << "\n";
		}
	}

	std::map<unsigned, unsigned>::reverse_iterator ii, end;
	ii = cluster_sizes.rbegin(); end = cluster_sizes.rend();
	ofs << "cluster size\tcluster count\n";
	ofs << "------------\t-------------\n";
	for (;ii != end; ++ii)
		ofs << ii->first << "\t" << ii->second << "\n";

	ofs << "\n\n";

	ofs << "cluster#\tcluster size\n";
	ofs << "--------\t------------\n";
	for (unsigned i = 0; i < clusters.size(); ++i)
	{
		unsigned cluster_size = clusters[i].size();
		ofs << i << "\t" << cluster_size << "\n";
	}


}

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//		Main

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
int main(int argc, char* argv[])
{
	t_portable_timer timer;

	// don't use C streams
	std::ios::sync_with_stdio(false);

	// saving parameters
	string exe_parameters;
	for (int i = 0; i < argc; ++i)
		exe_parameters += string(argv[i]) + ' ';

	// we use cout to print progress so flush per call
	cout.setf(std::ios_base::unitbuf);

	// get command line arguments
	t_program_args args;
	try
	{
	if (!process_command_line_options(argc, argv, args))
		return 1;

	//
	//	protein to gene id lookup
	//
	deque<unsigned>						prot_id_to_gene_id;
	t_index_string						prot_ids;
	t_index_string						gene_ids;
	{
		args.VERBOSE() << "\tReading gene to protein mapping...\n";
		t_progress_indicator 				dots(args.VERBOSE(),  5000);
		parse_protid_gene_id_from_file(	args.prot_to_gene_id,
										prot_id_to_gene_id,
										prot_ids,
										gene_ids,
										dots,
										args);
		dots.finish();
		args.VERBOSE() << "\t" << gene_ids.size() << " gene identifiers in total\n";
		args.VERBOSE() << "\t" << prot_ids.size() << " protein transcript identifiers in total\n";
	}


	//
	//	parse blast pairs
	//
	//
	deque<t_id_pair>			gene_pairs;
	deque<t_clustered_id_pair>	prot_pairs;
	unsigned 					cnt_pairs = 0;
	{
		unsigned DOT_SIZE = 5000;
		args.VERBOSE() << "\tReading homologue pairs from " << args.gene_prot_pairs << "\n";
		args.VERBOSE() << "\t\t(" << DOT_SIZE * 50 << " homologue pairs per line...)\n";
		t_progress_indicator 				dots(args.VERBOSE(),  DOT_SIZE);
		dots.use_timer();

		parse_blast_pairs_from_file(	prot_id_to_gene_id,
										prot_ids,
										gene_ids,
										gene_pairs,
										prot_pairs,
										args.gene_prot_pairs,
										dots);

		cnt_pairs = dots.value();
	}

	vector<unsigned>	gene_ids_to_cluster_lookup(gene_ids.size());
	args.VERBOSE() << "\tSingle linkage clustering genes\n";
	{
		unsigned DOT_SIZE = 50000;
		std::deque<std::vector<unsigned> >				clusters;
		args.VERBOSE() << "\t\t(" << DOT_SIZE * 50 << " homologue pairs per line...)\n";
		t_progress_indicator dots(args.VERBOSE(), DOT_SIZE);
		single_linkage_clusters(gene_pairs, clusters, dots);
		dots.finish();
		create_gene_ids_to_cluster_lookup(clusters, gene_ids_to_cluster_lookup);
		args.VERBOSE() << "\t" << clusters.size() << " clusters in total\n";

		// output sizes
		print_cluster_sizes(clusters, gene_ids, args);

	}

	// throw away pairs now to save memory
	deque<t_id_pair>().swap(gene_pairs);

	// annotate prot pairs with clusters
	assign_prot_pairs_to_clusters(prot_pairs, prot_id_to_gene_id, gene_ids_to_cluster_lookup);

	// throw away prot_id_to_gene_id now to save memory
	deque<unsigned>().swap(prot_id_to_gene_id);
	vector<unsigned>().swap(gene_ids_to_cluster_lookup);

	args.VERBOSE() << "\tSorting " << prot_pairs.size() << " protein pairs by clusters \n";
	sort(prot_pairs.begin(), prot_pairs.end());

	//
	// open results file
	//
	ofstream ofs(args.prot_pairs.c_str());
	if (!ofs)
		throw std::runtime_error("Could not open the results file to hold protein pairs [" + args.prot_pairs + "]");
	for (unsigned i = 0; i < prot_pairs.size(); ++i)
	{
		ofs << prot_ids.string_at(prot_pairs[i].first) << "\t" << prot_ids.string_at(prot_pairs[i].second) << "\n";
	}

	args.VERBOSE() << "\t" << string(60, '_') << "\n";
	args.VERBOSE() << "\t" << gene_ids.size() << " genes in total\n";
	args.VERBOSE() << "\t" << timer.reset() << " Seconds for clustering " << cnt_pairs << " blast pairs.\n";
	}
	catch (std::exception& e)
	{
		std_print_error(args.CERR(), e.what());
		return 1;
	}

	return 0;
}

